// 15_thread_pool.rs
// Rust 线程池设计教程

use std::thread;
use std::sync::{mpsc, Arc, Mutex};
use std::sync::mpsc::Receiver;

// 定义任务类型
type Job = Box<dyn FnOnce() + Send + 'static>;

// 线程池结构体
pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<mpsc::Sender<Job>>,
}

// 工作线程结构体
struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl ThreadPool {
    /// 创建指定大小的线程池
    /// 
    /// # 参数
    /// 
    /// * `size` - 线程池中线程的数量
    /// 
    /// # 返回值
    /// 
    /// 返回一个新的 ThreadPool 实例
    /// 
    /// # 示例
    /// 
    /// ```
    /// let pool = ThreadPool::new(4);
    /// ```
    pub fn new(size: usize) -> ThreadPool {
        assert!(size > 0);

        let (sender, receiver) = mpsc::channel();
        let receiver = Arc::new(Mutex::new(receiver));
        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        ThreadPool {
            workers,
            sender: Some(sender),
        }
    }

    /// 向线程池提交任务
    /// 
    /// # 参数
    /// 
    /// * `f` - 要执行的闭包任务
    /// 
    /// # 返回值
    /// 
    /// 如果成功提交任务则返回 Ok(())，否则返回 Err(())
    /// 
    /// # 示例
    /// 
    /// ```
    /// let pool = ThreadPool::new(4);
    /// pool.execute(|| {
    ///     println!("执行任务");
    /// });
    /// ```
    pub fn execute<F>(&self, f: F) -> Result<(), ()>
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);

        match &self.sender {
            Some(sender) => sender.send(job).map_err(|_| ()),
            None => Err(()),
        }
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        // 关闭发送端
        drop(self.sender.take());

        // 等待所有线程完成
        for worker in &mut self.workers {
            println!("关闭工作线程 {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
    }
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<Receiver<Job>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv();

            match message {
                Ok(job) => {
                    println!("工作线程 {} 获取任务", id);
                    job();
                }
                Err(_) => {
                    println!("工作线程 {} 断开连接", id);
                    break;
                }
            }
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}

fn main() {
    println!("=== 线程池设计教程 ===");

    // 1. 创建线程池
    println!("\n1. 创建线程池");
    let pool = ThreadPool::new(4);
    
    // 2. 提交任务到线程池
    println!("\n2. 提交任务到线程池");
    
    for i in 0..6 {
        let result = pool.execute(move || {
            println!("任务 {} 开始执行", i);
            thread::sleep(std::time::Duration::from_secs(1));
            println!("任务 {} 执行完成", i);
        });
        
        if result.is_ok() {
            println!("任务 {} 提交成功", i);
        } else {
            println!("任务 {} 提交失败", i);
        }
    }
    
    // 3. 提交需要共享数据的任务
    println!("\n3. 提交需要共享数据的任务");
    
    // 使用 Arc 共享可变数据
    let counter = Arc::new(Mutex::new(0));
    
    for _i in 0..10 {
        let counter = Arc::clone(&counter);
        let result = pool.execute(move || {
            let mut num = counter.lock().unwrap();
            *num += 1;
            println!("计数器值: {}", *num);
        });
        
        if result.is_err() {
            println!("任务提交失败");
        }
    }
    
    // 等待一段时间让所有任务完成
    thread::sleep(std::time::Duration::from_secs(3));
    
    println!("最终计数器值: {}", *counter.lock().unwrap());
    
    // 4. 线程池会自动清理
    println!("\n4. 程序结束，线程池自动清理");
    println!("线程池教程完成！");
}